package whoever;

import java.io.*;
import java.math.BigInteger;
import java.security.SecureRandom;
import java.util.Arrays;

import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.math.ec.ECCurve;
import org.bouncycastle.math.ec.ECPoint;

/**
 * SM2公钥加密算法实现 包括 -签名,验签 -密钥交换 -公钥加密 -私钥解密
 */
public class SM2 {
    // SM2系统参数作为公共参数，可以在所有的平台和系统之间公开共享。
    // 使用推荐曲线参数
    // 参数n,基点G的阶（n也是椭圆曲线的阶的素因子）
    private static BigInteger n = new BigInteger(
            "FFFFFFFE" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "7203DF6B" + "21C6052B" + "53BBF409" + "39D54123", 16);
    // 域的规模q = p
    private static BigInteger p = new BigInteger(
            "FFFFFFFE" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "00000000" + "FFFFFFFF" + "FFFFFFFF", 16);
    //a,b为Fq中元素，定义Fq上的一条椭圆曲线E
    private static BigInteger a = new BigInteger(
            "FFFFFFFE" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "00000000" + "FFFFFFFF" + "FFFFFFFC", 16);
    private static BigInteger b = new BigInteger(
            "28E9FA9E" + "9D9F5E34" + "4D5A9E4B" + "CF6509A7" + "F39789F5" + "15AB8F92" + "DDBCBD41" + "4D940E93", 16);
    // 基点G的坐标(gx,gy)
    private static BigInteger gx = new BigInteger(
            "32C4AE2C" + "1F198119" + "5F990446" + "6A39C994" + "8FE30BBF" + "F2660BE1" + "715A4589" + "334C74C7", 16);
    private static BigInteger gy = new BigInteger(
            "BC3736A2" + "F4F6779C" + "59BDCEE3" + "6B692153" + "D0A9877C" + "C62A4740" + "02DF32E5" + "2139F0A0", 16);
    private static ECDomainParameters ecc_bc_spec;
    // 密钥交换协议用到的参数w和_2w
    private static int w = (int) Math.ceil(n.bitLength() * 1.0 / 2) - 1;
    private static BigInteger _2w = new BigInteger("2").pow(w);
    // 32字节，256位
    private static final int DIGEST_LENGTH = 32;
    // 随机数生成器
    private static SecureRandom random = new SecureRandom();
    // ECCurve是抽象类，使用其嵌套类ECCurve.Fp
    private static ECCurve.Fp curve;
    // 基点G
    private static ECPoint G;
    private boolean debug = false;

    // 构造类
    SM2() {
        curve = new ECCurve.Fp(p, // q
                a, // a
                b); // b
        G = curve.createPoint(gx, gy);
        ecc_bc_spec = new ECDomainParameters(curve, G, n);
    }

    public SM2(boolean debug) {
        this();
        this.debug = debug;
    }

    /**
     * 测试用：以16进制字符串显示公钥地址; byte[]->16进制的String
     */
    static void printHexString(byte[] b) {
        for (byte value : b) {
            // 保持byte-->integer二进制补码的一致性
            String hex = Integer.toHexString(value & 0xFF);
            if (hex.length() == 1) {
                hex = '0' + hex;
            }
            System.out.print(hex.toUpperCase());
        }
        System.out.println();
    }

    /**
     * 给定上界的随机数生成器
     * @param max 规定随机数上界 BigInteger
     * @return 随机数
     */
    private static BigInteger random(BigInteger max) {
        // 生成随机数数r, 0 <= r <= 2^256-1
        // 此处的random参数是SecureRandom类
        BigInteger r = new BigInteger(256, random);
        // 若r大于或等于上界
        while (r.compareTo(max) >= 0) {
            r = new BigInteger(128, random);
        }
        return r;
    }

    /**
     * 判断字节数组是否全0
     * @param buffer 传入字符数组 byte[]
     * @return 全零为true，否则为false
     */
    private boolean allZero(byte[] buffer) {
        for (byte value : buffer) {
            if (value != 0)
                return false;
        }
        return true;
    }

    /**
     * 判断生成参数是否在范围内
     * @param param 参数 BigInteger
     * @param min 下界 BigInteger
     * @param max 上界 BigInterger
     * @return 判断结果 bool
     */
    private boolean between(BigInteger param, BigInteger min, BigInteger max) {
        return param.compareTo(min) >= 0 && param.compareTo(max) < 0;
    }

    /**
     * 判断生成的公钥是否合法
     * @param publicKey 公钥 ECPoint
     * @return 是否合法
     */
    private boolean checkPublicKey(ECPoint publicKey) {

        if (!publicKey.isInfinity()) {

            BigInteger x = publicKey.getXCoord().toBigInteger();
            BigInteger y = publicKey.getYCoord().toBigInteger();

            if (between(x, new BigInteger("0"), p) && between(y, new BigInteger("0"), p)) {

                BigInteger xResult = x.pow(3).add(a.multiply(x)).add(b).mod(p);

                if (debug)
                    System.out.println("xResult: " + xResult.toString());

                BigInteger yResult = y.pow(2).mod(p);

                if (debug)
                    System.out.println("yResult: " + yResult.toString());

                return yResult.equals(xResult) && publicKey.multiply(n).isInfinity();
            }
        }
        return false;
    }

    /**
     * 密钥派生函数KDF：从一个共享的秘密比特串中派生出密钥数据
     * @param Z 比特串 byte[]
     * @param klen 要获得的密钥数据的长度 int
     * @return 长度为klen的密钥数据比特串
     */
    private static byte[] KDF(byte[] Z, int klen) {
        // a)初始化计数器
        int ct = 1;
        // b)迭代计算Hai
        int end = (int) Math.ceil(klen * 1.0 / 32);
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        try {
            for (int i = 1; i < end; i++) {
                baos.write(sm3hash(Z, SM3.toByteArray(ct)));
                ct++;
            }
            byte[] last = sm3hash(Z, SM3.toByteArray(ct));
            // c)若 klen/32 为整数
            // 注：klen为字节长度，计算迭代次数时实际为klen * 8 / 256,因为sm3输出256位杂凑值
            if (klen % 32 == 0) {
                baos.write(last);
            }
            // c)若 keln/32 不为整数
            else
                baos.write(last, 0, klen % 32);
            // d)拼接返回
            return baos.toByteArray();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 生成密钥对
     * @return 密钥对 SM2KeyPair
     */
    public SM2KeyPair generateKeyPair() {
        // 使用当前SM2类的随机数发生器
        BigInteger d = random(n.subtract(new BigInteger("1")));

        // G倍乘随机数后归一化作为公钥ECPoint
        // 归一化保证射影坐标为1，因此(x,y)反应射影坐标系的等效点
        SM2KeyPair keyPair = new SM2KeyPair(G.multiply(d).normalize(), d);

        if (checkPublicKey(keyPair.getPublicKey())) {
            if (debug)
                System.out.println("generate key successfully");
            return keyPair;
        } else {
            if (debug)
                System.err.println("generate key failed");
            return null;
        }
    }

    /**
     * 公钥加密
     * @param input 明文 String
     * @param publicKey 公钥 ECPoint
     * @return 密文 byte[]
     */
    public byte[] encrypt( String input, ECPoint publicKey) {
        byte[] inputBuffer = input.getBytes();

        if (debug)
            printHexString(inputBuffer);

        byte[] C1Buffer;
        ECPoint kpb;
        byte[] t;
        do {
            // A1 产生随机数k，k属于[1, n-1]
            BigInteger k = random(n);
            if (debug) {
                System.out.print("k: ");
                printHexString(k.toByteArray());
            }

            // A2 倍点运算，计算椭圆曲线点C1 = [k]G = (x1, y1)
            ECPoint C1 = G.multiply(k);
            // 将C1转化为比特串，不进行点压缩
            C1Buffer = C1.getEncoded(false);
            if (debug) {
                System.out.print("C1: ");
                printHexString(C1Buffer);
            }

            // A3 计算椭圆曲线点S = [h]Pb = (x2, y2)
            BigInteger h = ecc_bc_spec.getH();
            if (h != null) {
                ECPoint S = publicKey.multiply(h);
                // 若S是无穷远点，报错退出
                if (S.isInfinity())
                    throw new IllegalStateException();
            }

            // A4 计算 [k]PB = (x2, y2),并转化为比特串
            kpb = publicKey.multiply(k).normalize();
            byte[] kpbBytes = kpb.getEncoded(false);

            // A5 计算 t = KDF(x2||y2, klen)
            t = KDF(kpbBytes, inputBuffer.length);
        } while (allZero(t));

        // A6 计算 C2 = M^t,异或运算
        byte[] C2 = new byte[inputBuffer.length];
        for (int i = 0; i < inputBuffer.length; i++) {
            C2[i] = (byte) (inputBuffer[i] ^ t[i]);
        }

        // A7 计算 C3 = Hash(x2 || M || y2)
        byte[] C3 = sm3hash(kpb.getXCoord().toBigInteger().toByteArray(), inputBuffer,
                kpb.getYCoord().toBigInteger().toByteArray());

        // A8 输出密文 C = C1 || C2 || C3
        byte[] encryptResult = new byte[C1Buffer.length + C2.length + C3.length];
        System.arraycopy(C1Buffer, 0, encryptResult, 0, C1Buffer.length);
        System.arraycopy(C2, 0, encryptResult, C1Buffer.length, C2.length);
        System.arraycopy(C3, 0, encryptResult, C1Buffer.length + C2.length, C3.length);

        if (debug) {
            System.out.print("密文: ");
            printHexString(encryptResult);
        }

        return encryptResult;
    }

    /**
     * 私钥解密
     * @param encryptData 密文 byte[]
     * @param privateKey 私钥 BigInteger
     * @return 明文 byte[]
     */
    public String decrypt(byte[] encryptData, BigInteger privateKey) {

        if (debug)
            System.out.println("encryptData length: " + encryptData.length);

        // B1 取出密文中的C1,转化为曲线上的点并验证是否满足椭圆方程
        byte[] C1Byte = new byte[65];
        System.arraycopy(encryptData, 0, C1Byte, 0, C1Byte.length);
        ECPoint C1 = curve.decodePoint(C1Byte).normalize();

        // B2 计算椭圆曲线点S = [h]C1 是否为无穷点
        BigInteger h = ecc_bc_spec.getH();
        if (h != null) {
            ECPoint S = C1.multiply(h);
            if (S.isInfinity())
                throw new IllegalStateException();
        }

        // B3 计算 [dB]C1 = (x2, y2)
        ECPoint dBC1 = C1.multiply(privateKey).normalize();

        // B4 计算 t = KDF(x2 || y2, klen)
        byte[] dBC1Bytes = dBC1.getEncoded(false);
        int klen = encryptData.length - 65 - DIGEST_LENGTH;
        byte[] t = KDF(dBC1Bytes, klen);

        // 若t为全0，报错退出
        if (allZero(t)) {
            System.err.println("all zero");
            throw new IllegalStateException();
        }

        // B5 去处c2，并计算 M'=C2^t （异或运算）
        byte[] M = new byte[klen];
        for (int i = 0; i < M.length; i++) {
            M[i] = (byte) (encryptData[C1Byte.length + i] ^ t[i]);
        }
        if (debug)
            printHexString(M);

        // B6 计算 u = Hash(x2 || M' || y2) 判断 u == C3是否成立
        byte[] C3 = new byte[DIGEST_LENGTH];

        if (debug)
            try {
                System.out.println("M = " + new String(M, "UTF8"));
            } catch (UnsupportedEncodingException e1) {
                e1.printStackTrace();
            }

        // 取出C3
        System.arraycopy(encryptData, encryptData.length - DIGEST_LENGTH, C3, 0, DIGEST_LENGTH);
        // 计算 Hash(x2 || M' || y2)
        byte[] u = sm3hash(dBC1.getXCoord().toBigInteger().toByteArray(), M,
                dBC1.getYCoord().toBigInteger().toByteArray());
        if (Arrays.equals(u, C3)) {
            if (debug)
                System.out.println("解密成功");
            try {
                return new String(M, "UTF8");
            } catch (UnsupportedEncodingException e) {
                e.printStackTrace();
            }
            return null;
        } else {
            if (debug) {
                System.out.print("u = ");
                printHexString(u);
                System.out.print("C3 = ");
                printHexString(C3);
                System.err.println("解密验证失败");
            }
            return null;
        }

    }

    /**
     * 导出公钥到本地
     * @param publicKey 待导出的公钥 ECPoint
     * @param path 储存公钥的文件路径
     */
    public void exportPublicKey(ECPoint publicKey, String path) {
        File file = new File(path);
        try {
            if (!file.exists())
                file.createNewFile();
            byte[] buffer = publicKey.getEncoded(false);
            FileOutputStream fos = new FileOutputStream(file);
            fos.write(buffer);
            fos.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 从本地导入公钥
     * @param path 储存公钥的文件路径 String
     * @return 公钥 ECPoint
     */
    public ECPoint importPublicKey(String path) {
        File file = new File(path);
        try {
            if (!file.exists())
                return null;
            FileInputStream fis = new FileInputStream(file);
            ByteArrayOutputStream baos = new ByteArrayOutputStream();

            byte[] buffer = new byte[16];
            int size;
            while ((size = fis.read(buffer)) != -1) {
                baos.write(buffer, 0, size);
            }
            fis.close();
            return curve.decodePoint(baos.toByteArray());
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 导出私钥到本地
     * 注：导出的是序列化的Integer类型
     * @param privateKey 待导出的私钥 BigInteger
     * @param path 储存私钥的文件路径
     */
    public void exportPrivateKey(BigInteger privateKey, String path) {
        File file = new File(path);
        try {
            if (!file.exists())
                file.createNewFile();
            ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(file));
            oos.writeObject(privateKey);
            oos.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 从本地导入私钥
     * 注：导入的是序列化对象
     * @param path 储存私钥的文件路径
     * @return 私钥 BigInteger
     */
    public BigInteger importPrivateKey(String path) {
        File file = new File(path);
        try {
            if (!file.exists())
                return null;
            FileInputStream fis = new FileInputStream(file);
            ObjectInputStream ois = new ObjectInputStream(fis);
            BigInteger res = (BigInteger) (ois.readObject());
            ois.close();
            fis.close();
            return res;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    // 公钥 ECPoint -> String
    public static String getPublicKeyInString(ECPoint publicKeyInEcPoint) {
        byte[] b = publicKeyInEcPoint.getEncoded(false);
        StringBuilder res = new StringBuilder();
        for (int i = 0; i < b.length; i++) {
            // 保持byte-->integer二进制补码的一致性
            String hex = Integer.toHexString(b[i] & 0xFF);
            if (hex.length() == 1) {
                hex = '0' + hex;
            }
            res.append(hex);
        }
        return res.toString();
    }

    // 公钥 String -> ECPoint
    private static byte charToByte(char c) {
        return (byte) "0123456789abcdef".indexOf(c);
    }
    public static ECPoint getPublicKeyInECPoint(String publicKeyInString) {
        // 16进制String->byte[]
        char[] hexChars = publicKeyInString.toCharArray();
        int length = publicKeyInString.length() / 2;
        byte[] b = new byte[length];
        for (int i = 0; i < length; i++) {
            int pos = i * 2;
            b[i] = (byte) (charToByte(hexChars[pos]) << 4 | charToByte(hexChars[pos + 1]));
        }
        return curve.decodePoint(b);
    }


    // 私钥 BigInteger -> String
    public static String getPrivateKeyInString(BigInteger privateKeyInBigInteger){
        return privateKeyInBigInteger.toString();
    }

    // 私钥  String -> BigInteger
    public static BigInteger getPrivateKeyInBigInteger(String privateKeyInString){
        BigInteger privateKeyInBigInteger = new BigInteger(privateKeyInString);
        return privateKeyInBigInteger;
    }

    /**
     * 字节数组拼接
     * @param params 若干字节数组 (byte1[], byte2[],...)
     * @return 连接后的字节数组 byte[]
     */
    private static byte[] join(byte[]... params) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        byte[] res = null;
        try {
            for (byte[] param : params) {
                baos.write(param);
            }
            res = baos.toByteArray();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return res;
    }

    /**
     * 使用sm3算法计算若干字符数组拼接后的杂凑值
     * @param params 若干字符数组 byte[]...
     * @return 杂凑值 byte[]
     */
    private static byte[] sm3hash(byte[]... params) {
        byte[] res = null;
        try {
            res = SM3.hash(join(params));
        } catch (IOException e) {
            e.printStackTrace();
        }
        return res;
    }

    /**
     * SM2签名中Signature类的构造
     */
    public static class Signature {
        BigInteger r;
        BigInteger s;

        public Signature(BigInteger r, BigInteger s) {
            this.r = r;
            this.s = s;
        }

        // 将(BigInteger r, BigInteger s)转化为(String r, String s)
        public String toString() {
            return r.toString(16) +
                    "," + s.toString(16);
        }
    }

    /**
     * 计算签名过程中的ZA值
     * @param IDA 用户标识 String
     * @param aPublicKey 公钥 ECPoint
     * @return ZA值 byte[]
     */
    private static byte[] ZA(String IDA, ECPoint aPublicKey) {
        byte[] idaBytes = IDA.getBytes();
        // entlenA为标识的比特长度
        int entlenA = idaBytes.length * 8;
        // ENTLA为比特长度整数转换而成的两个字节
        // 如长度为144bits = Ox90bits,字节为(00,90)
        byte[] ENTLA = new byte[] { (byte) (entlenA & 0xFF00), (byte) (entlenA & 0x00FF) };
        // 计算ZA值
        byte[] ZA = sm3hash(ENTLA, idaBytes,
                a.toByteArray(), b.toByteArray(),
                gx.toByteArray(), gy.toByteArray(),
                aPublicKey.getXCoord().toBigInteger().toByteArray(),
                aPublicKey.getYCoord().toBigInteger().toByteArray());
        return ZA;
    }

    /**
     * 签名者签名过程
     * @param M 签名信息 String
     * @param IDA 签名方标识 String
     * @param keyPair 签名方公钥私钥密钥对 SM2KeyPair
     * @return 签名 Signature
     */
    public Signature sign(String M, String IDA, SM2KeyPair keyPair) {
        // A1过程
        byte[] ZA = ZA(IDA, keyPair.getPublicKey());
        // A2过程
        byte[] M_ = join(ZA, M.getBytes());
        // 把byte[]类型的M_转化为BigInteger
        BigInteger e = new BigInteger(1, sm3hash(M_));
        BigInteger k;
        BigInteger r;
        BigInteger s;
        do {
            do {
            // A3过程
            k = random(n);
            // A4过程
            ECPoint p1 = G.multiply(k).normalize();
            BigInteger x1 = p1.getXCoord().toBigInteger();
            // A5过程
            r = e.add(x1);
            r = r.mod(n);
        } while (r.equals(BigInteger.ZERO) || r.add(k).equals(n));
        // A6过程
            s = ((keyPair.getPrivateKey().add(BigInteger.ONE).modInverse(n))
            .multiply((k.subtract(r.multiply(keyPair.getPrivateKey()))).mod(n))).mod(n);
        }while (s.equals(BigInteger.ZERO));

        return new Signature(r, s);
    }

    /**
     * 验证方验签
     * @param M 签名信息 String
     * @param signature 签名 Signature
     * @param IDA 签名方标识 String
     * @param aPublicKey 签名方公钥 ECPoint
     * @return 结果 true or false
     */
    public boolean verify(String M, Signature signature, String IDA, ECPoint aPublicKey) {
        // B1过程
        if (!between(signature.r, BigInteger.ONE, n))
            return false;
        // B2过程
        if (!between(signature.s, BigInteger.ONE, n))
            return false;
        // B3过程
        byte[] M_ = join(ZA(IDA, aPublicKey), M.getBytes());
        // B4过程
        BigInteger e = new BigInteger(1, sm3hash(M_));
        // B5过程
        BigInteger t = signature.r.add(signature.s).mod(n);
        if (t.equals(BigInteger.ZERO))
            return false;
        // B6过程
        ECPoint p1 = G.multiply(signature.s).normalize();
        ECPoint p2 = aPublicKey.multiply(t).normalize();
        BigInteger x1 = p1.add(p2).normalize().getXCoord().toBigInteger();
        // B7过程
        BigInteger R = e.add(x1).mod(n);
        if (R.equals(signature.r))
            return true;
        return false;
    }

    /**
     * 传输实体类，用于密钥交换协议
     */
    private static class TransportEntity implements Serializable {
        final byte[] R; //R点
        final byte[] S; //验证S
        final byte[] Z; //用户标识
        final byte[] K; //公钥

        // 构造类
        public TransportEntity(byte[] r, byte[] s, byte[] z, ECPoint pKey) {
            R = r;
            S = s;
            Z=z;
            K=pKey.getEncoded(false);
        }
    }

    /**
     * 密钥协商、交换过程
     * 设A、B用户进行协商密钥，A为发起方，B为响应方
     */
    public static class KeyExchange {
        BigInteger rA;
        ECPoint RA;
        ECPoint V;
        byte[] Z;
        byte[] key;

        String ID;
        SM2KeyPair keyPair;

        // 构造类
        public KeyExchange(String ID, SM2KeyPair keyPair) {
            this.ID=ID;
            this.keyPair = keyPair;
            this.Z=ZA(ID, keyPair.getPublicKey());
        }

        /**
         * 密钥协商第一步，由A发起
         * @return 传输实体 TransportEntity
         */
        public TransportEntity keyExchange_1() {
            // A1 生成随机数rA属于[1,n-1]
            rA = random(n);
            // A2 计算椭圆曲线点RA
            RA = G.multiply(rA).normalize();
            // A3 将RA发送给B
            return new TransportEntity(RA.getEncoded(false), null,Z,keyPair.getPublicKey());
        }

        /**
         * 密钥协商第二步，B做出响应
         * @param entity A第一步的传输实体 TransportEntity
         * @return 传输实体 TransportEntity
         */
        public TransportEntity keyExchange_2(TransportEntity entity) {
            // B1 生成随机数rB属于[1,n-1]
            BigInteger rB = random(n);
            // B2 计算椭圆曲线点RB
            ECPoint RB = G.multiply(rB).normalize();

            this.rA=rB;
            this.RA=RB;

            // B3 从RB中取出x2并计算_x2
            BigInteger x2 = RB.getXCoord().toBigInteger();
            x2 = _2w.add(x2.and(_2w.subtract(BigInteger.ONE)));

            // B4 计算tB
            BigInteger tB = keyPair.getPrivateKey().add(x2.multiply(rB)).mod(n);

            // B5 计算_x1
            ECPoint RA = curve.decodePoint(entity.R).normalize();
            BigInteger x1 = RA.getXCoord().toBigInteger();
            x1 = _2w.add(x1.and(_2w.subtract(BigInteger.ONE)));

            // B6 计算点V
            ECPoint aPublicKey=curve.decodePoint(entity.K).normalize();
            ECPoint temp = aPublicKey.add(RA.multiply(x1).normalize()).normalize();
            ECPoint V = temp.multiply(ecc_bc_spec.getH().multiply(tB)).normalize();
            if (V.isInfinity())
                throw new IllegalStateException();
            this.V=V;

            byte[] xV = V.getXCoord().toBigInteger().toByteArray();
            byte[] yV = V.getYCoord().toBigInteger().toByteArray();

            // B7 计算KB
            byte[] KB = KDF(join(xV, yV, entity.Z, this.Z), 16);
            key = KB;
            System.out.print("协商得B密钥:");
            printHexString(KB);
            // B8 计算SB
            byte[] sB = sm3hash(new byte[] { 0x02 }, yV,
                    sm3hash(xV, entity.Z, this.Z, RA.getXCoord().toBigInteger().toByteArray(),
                            RA.getYCoord().toBigInteger().toByteArray(), RB.getXCoord().toBigInteger().toByteArray(),
                            RB.getYCoord().toBigInteger().toByteArray()));
            // B9 发送RB、SB
            return new TransportEntity(RB.getEncoded(false), sB,this.Z,keyPair.getPublicKey());
        }

        /**
         * 密钥协商第三步，由A发起
         * @param entity B第二步的传输实体
         * @return 传输实体 TransportEntity
         */
        public TransportEntity keyExchange_3(TransportEntity entity) {
            // A4 从RA取出x1，并计算_x1
            BigInteger x1 = RA.getXCoord().toBigInteger();
            x1 = _2w.add(x1.and(_2w.subtract(BigInteger.ONE)));

            // A5 计算tA
            BigInteger tA = keyPair.getPrivateKey().add(x1.multiply(rA)).mod(n);
            ECPoint RB = curve.decodePoint(entity.R).normalize();

            // A6 取出x2，并计算_x2
            BigInteger x2 = RB.getXCoord().toBigInteger();
            x2 = _2w.add(x2.and(_2w.subtract(BigInteger.ONE)));

            // A7 计算点U
            ECPoint bPublicKey=curve.decodePoint(entity.K).normalize();
            ECPoint temp = bPublicKey.add(RB.multiply(x2).normalize()).normalize();
            ECPoint U = temp.multiply(ecc_bc_spec.getH().multiply(tA)).normalize();
            if (U.isInfinity())
                throw new IllegalStateException();
            this.V=U;

            // A8 计算KA
            byte[] xU = U.getXCoord().toBigInteger().toByteArray();
            byte[] yU = U.getYCoord().toBigInteger().toByteArray();
            byte[] KA = KDF(join(xU, yU,
                    this.Z, entity.Z), 16);
            key = KA;
            System.out.print("协商得A密钥:");
            printHexString(KA);
            // A9 计算S1
            byte[] s1= sm3hash(new byte[] { 0x02 }, yU,
                    sm3hash(xU, this.Z, entity.Z, RA.getXCoord().toBigInteger().toByteArray(),
                            RA.getYCoord().toBigInteger().toByteArray(), RB.getXCoord().toBigInteger().toByteArray(),
                            RB.getYCoord().toBigInteger().toByteArray()));
            if(Arrays.equals(entity.S, s1))
                System.out.println("B->A 密钥确认成功");
            else
                System.out.println("B->A 密钥确认失败");

            // A10 计算SA并发送给B
            byte[] sA= sm3hash(new byte[] { 0x03 }, yU,
                    sm3hash(xU, this.Z, entity.Z, RA.getXCoord().toBigInteger().toByteArray(),
                            RA.getYCoord().toBigInteger().toByteArray(), RB.getXCoord().toBigInteger().toByteArray(),
                            RB.getYCoord().toBigInteger().toByteArray()));

            return new TransportEntity(RA.getEncoded(false), sA,this.Z,keyPair.getPublicKey());
        }

        /**
         * 密钥确认第四步也是最后一步，由用户B完成
         * @param entity A第三步的传输实体
         */
        public void keyExchange_4(TransportEntity entity) {
            // 计算 SA == S2,若相等，完成传输
            byte[] xV = V.getXCoord().toBigInteger().toByteArray();
            byte[] yV = V.getYCoord().toBigInteger().toByteArray();
            ECPoint RA = curve.decodePoint(entity.R).normalize();
            byte[] s2= sm3hash(new byte[] { 0x03 }, yV,
                    sm3hash(xV, entity.Z, this.Z, RA.getXCoord().toBigInteger().toByteArray(),
                            RA.getYCoord().toBigInteger().toByteArray(), this.RA.getXCoord().toBigInteger().toByteArray(),
                            this.RA.getYCoord().toBigInteger().toByteArray()));
            if(Arrays.equals(entity.S, s2))
                System.out.println("A->B 密钥确认成功");
            else
                System.out.println("A->B 密钥确认失败");
        }
    }


    public boolean isDebug() {
        return debug;
    }

    public void setDebug(boolean debug) {
        this.debug = debug;
    }

    // 测试密钥协商
    public static void main(String[] args) throws UnsupportedEncodingException {

        SM2 sm02 = new SM2();

        System.out.println("-----------------密钥协商-----------------");
        String aID = "AAAAAAAAAAAAA";
        SM2KeyPair aKeyPair = sm02.generateKeyPair();
        KeyExchange aKeyExchange = new KeyExchange(aID,aKeyPair);

        String bID = "BBBBBBBBBBBBB";
        SM2KeyPair bKeyPair = sm02.generateKeyPair();
        KeyExchange bKeyExchange = new KeyExchange(bID,bKeyPair);
        TransportEntity entity1 = aKeyExchange.keyExchange_1();
        TransportEntity entity2 = bKeyExchange.keyExchange_2(entity1);
        TransportEntity entity3 = aKeyExchange.keyExchange_3(entity2);
        bKeyExchange.keyExchange_4(entity3);
    }

}

